# coding=utf-8

import csv
import sys
import os

'''
Type,Name,Request Count,Failure Count,Median Response Time,Average Response Time,Min Response Time,Max Response Time,Average Content Size,Requests/s,Failures/s,50%,66%,75%,80%,90%,95%,98%,99%,99.90%,99.99%,100%
trpc,error,2,2,190,195.40250702993944,186.2806520657614,204.5243619941175,0,0.001700541,0.001700541,200,200,200,200,200,200,200,200,200,200,200
trpc,first,59,0,300,471.12733045379804,152.64103200752288,2150.8804330369458,3663.2033898305085,0.050165955,0,300,440,490,740,1100,1400,1800,2200,2200,2200,2200
trpc,token_per_second,59,0,50,49.11469637014461,34.36401392,58.152598062952414,0,0.050165955,0,50,51,52,53,56,57,58,58,58,58,58
trpc,whole,59,0,17000,19919.127817918798,889.6602810127661,100965.14288592152,945.4406779661017,0.050165955,0,17000,25000,29000,30000,35000,42000,51000,101000,101000,101000,101000
,Aggregated,179,2,280,6739.182267,34.36401392,100965.14288592152,1519.0502793296089,0.1521984048906079,0.001700541,280,1200,7800,15000,27000,33000,41000,51000,101000,101000,101000


    {'Type': 'trpc', 'Name': 'error', 'Request Count': '2', 'Failure Count': '2', 'Median Response Time': '190',
     'Average Response Time': '195.40250702993944', 'Min Response Time': '186.2806520657614',
     'Max Response Time': '204.5243619941175', 'Average Content Size': '0', 'Requests/s': '0.001700541',
     'Failures/s': '0.001700541', '50%': '200', '66%': '200', '75%': '200', '80%': '200', '90%': '200', '95%': '200',
     '98%': '200', '99%': '200', '99.90%': '200', '99.99%': '200', '100%': '200'}
'''

def deal_data(data):
    tmp = {}
    tmp["SUC_REQUEST_COUNT"] = int(data['whole']['Request Count'])
    print(data)
    if 'error' in data:
        tmp["FAIL_REQUEST_COUNT"] = int(data['error']['Failure Count'])
    else:
        tmp["FAIL_REQUEST_COUNT"] = 0
    tmp["REQUEST_COUNT"] = tmp["SUC_REQUEST_COUNT"] + tmp["FAIL_REQUEST_COUNT"]
    tmp['FAIL_RATE'] = 100 * tmp['FAIL_REQUEST_COUNT'] / tmp['REQUEST_COUNT']


    # qps qpm
    tmp['QPS'] = float(data['whole']['Requests/s'])
    tmp['QPM'] = tmp['QPS'] * 60


    # 首token耗时
    tmp['FIRST_TOKEN_TIME'] = float(data['first']['Average Response Time'])
    tmp['FIRST_TOKEN_TIME_P50'] = float(data['first']['50%'])
    tmp['FIRST_TOKEN_TIME_P90'] = float(data['first']['90%'])
    tmp['FIRST_TOKEN_TIME_P95'] = float(data['first']['95%'])
    tmp['FIRST_TOKEN_TIME_P99'] = float(data['first']['99%'])

    # input token数量
    tmp['INPUT_TOKEN'] = float(data['input_token']['Average Response Time'])

    # output token数量
    tmp['OUTPUT_TOKEN'] = float(data['output_token']['Average Response Time'])

    # 请求耗时
    tmp['WHOLE_TOKEN_TIME'] = float(data['whole']['Average Response Time'])
    tmp['WHOLE_TOKEN_TIME_P50'] = float(data['whole']['50%'])
    tmp['WHOLE_TOKEN_TIME_P90'] = float(data['whole']['90%'])
    tmp['WHOLE_TOKEN_TIME_P95'] = float(data['whole']['95%'])
    tmp['WHOLE_TOKEN_TIME_P99'] = float(data['whole']['99%'])

    # 平均输入输出token处理耗时
    tmp['OUTPUT_TIME'] = tmp['WHOLE_TOKEN_TIME'] - tmp['FIRST_TOKEN_TIME']
    tmp['IPNUT_TIME'] = tmp['FIRST_TOKEN_TIME']

    # # 输入输出token
    # tmp['INPUT_TOKEN'] = float(data['first']['Average Content Size'])
    # tmp['OUTPUT_TOKEN'] = float(data['whole']['Average Content Size'])

    # # 每并发每秒输入token数
    # tmp['INPUT_TOKEN_PER_SECOND_CONCURRENCY'] = tmp['INPUT_TOKEN'] / tmp['IPNUT_TIME'] * 1000
    #
    # # 每并发每秒输出token数
    # tmp['OUTPUT_TOKEN_PER_SECOND_CONCURRENCY'] = tmp['OUTPUT_TOKEN'] / tmp['OUTPUT_TIME'] * 1000

    # 平均每秒输出token数
    tmp['OUTPUT_TOKEN_PER_SECOND'] = float(data['token_per_second']['Average Response Time'])
    tmp['OUTPUT_TOKEN_PER_SECOND_P50'] = float(data['whole']['50%'])
    tmp['OUTPUT_TOKEN_PER_SECOND_P90'] = float(data['whole']['90%'])
    tmp['OUTPUT_TOKEN_PER_SECOND_P95'] = float(data['whole']['95%'])
    tmp['OUTPUT_TOKEN_PER_SECOND_P99'] = float(data['whole']['99%'])

    return tmp


def main():
    if len(sys.argv) < 2:
        print(f"Error: No filename provided")
        print(f"Usage: python3 {sys.argv[0]} <filename>")
        sys.exit(1)

    filename = sys.argv[1]
    basename = os.path.splitext(os.path.basename(filename))[0]
    print(f"filename {filename} basename {basename}")
    with open(filename, mode='r', encoding='utf-8') as file:
        # 创建一个CSV阅读器
        csv_reader = csv.DictReader(file)

        ret = []
        # 遍历每一行
        data = {}
        for row in csv_reader:
            # 解析每一行的数据
            if row['Type'] == '':
                print(data.keys())
                ret.append(deal_data(data))
            elif row['Type'] == 'Type':
                data = {}
                continue
            data[row['Name']] = row


    # headers = ['QPM', 'QPS', 'INPUT_TOKEN', 'OUTPUT_TOKEN',
    #            'FIRST_TOKEN_TIME', 'FIRST_TOKEN_TIME_P50', 'FIRST_TOKEN_TIME_P90', 'FIRST_TOKEN_TIME_P99',
    #            'WHOLE_TOKEN_TIME', 'WHOLE_TOKEN_TIME_P50', 'WHOLE_TOKEN_TIME_P90', 'WHOLE_TOKEN_TIME_P99',
    #            'OUTPUT_TIME', 'INPUT_TOKEN_PER_SECOND_CONCURRENCY',
    #            'OUTPUT_TOKEN_PER_SECOND_CONCURRENCY', 'OUTPUT_TOKEN_PER_SECOND',
    #            'SUC_REQUEST_COUNT', 'FAIL_REQUEST_COUNT', 'FAIL_RATE',]

    headers = ['QPM', 'QPS', 'INPUT_TOKEN', 'OUTPUT_TOKEN',
               'FIRST_TOKEN_TIME', 'FIRST_TOKEN_TIME_P50', 'FIRST_TOKEN_TIME_P90', 'FIRST_TOKEN_TIME_P99',
               'WHOLE_TOKEN_TIME', 'WHOLE_TOKEN_TIME_P50', 'WHOLE_TOKEN_TIME_P90', 'WHOLE_TOKEN_TIME_P99',
               'OUTPUT_TIME',  'OUTPUT_TOKEN_PER_SECOND',
               'SUC_REQUEST_COUNT', 'FAIL_REQUEST_COUNT', 'FAIL_RATE', ]

    # ret按header顺序写入csv
    with open(f'result-{basename}.csv', mode='w', encoding='utf-8') as file:
        csv_writer = csv.DictWriter(file, fieldnames=headers)
        csv_writer.writeheader()
        for row in ret:
            filtered_row = {key: row[key] for key in headers if key in row}
            csv_writer.writerow(filtered_row)

    print(f"result-{basename}.csv saved")

if __name__ == '__main__':
    main()
